# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import functools
import logging

import pandas as pd

from coin.base.datetime_util import to_datetime, to_timestamp_int
from experimental.prophet.ops.util import to_duration_int, to_duration_str

from experimental.prophet.graph import graph
from experimental.prophet.graph_runner.fastfeed import run_from_fastfeed_multiprocess
from experimental.prophet.craps.universe import universe_btc

from experimental.prophet.ops import globals
from experimental.prophet.ops.aggregator import scalar_value_aggregator
from experimental.prophet.ops.dataframe_feeder import dataframe_feeder
from experimental.prophet.ops.debug import printer
from experimental.prophet.ops.fastfeed import fastfeed_coin
from experimental.prophet.ops import timeseries
from experimental.prophet.ops import y_gen

DEFAULT_TIMED_Y_POINTS = [
    '100ms', '500ms', '1s', '2s', '5s', '10s', '30s', '1m', '5m', '10m', '30m', '1h'
]


def get_symbol_info(symbol, ref_ts):
  universe = universe_btc(ref_ts)
  for sinfo in universe:
    if symbol == str(sinfo.product):
      return sinfo
  raise KeyError(symbol)


def get_date_timestamp_range(date):
  from_ts = to_datetime(date)
  to_ts = from_ts + datetime.timedelta(days=1)
  return to_timestamp_int(from_ts), to_timestamp_int(to_ts)


def timed_y(symbol_info, machine, fills, dump_from_ts, dump_to_ts, windows):
  ts = globals.global_timestamp()

  with graph.namespace('timed_y'):
    dump_from_ts = to_timestamp_int(dump_from_ts)
    dump_to_ts = to_timestamp_int(dump_to_ts)

    feed = fastfeed_coin(symbol_info.product, symbol_info.sub_req, machine)
    book = graph.create_node('fastfeed.BookBestLevel', [feed.book])
    ask0 = book.outputs[1].named_ref('ask0')
    bid0 = book.outputs[3].named_ref('bid0')
    mid = ((ask0 + bid0) / 2.).named_ref('mid')

    ask0_past = reversed(timeseries.time_shift(ask0, windows))
    bid0_past = reversed(timeseries.time_shift(bid0, windows))
    mid_past = reversed(timeseries.time_shift(mid, windows))

    ask0_fut = y_gen.time_shift_y(ask0, windows)
    bid0_fut = y_gen.time_shift_y(bid0, windows)
    mid_fut = y_gen.time_shift_y(mid, windows)

    # fill
    df = pd.DataFrame({
        'timestamp': [f.ts for f in fills],
        'fill_price': [f.price for f in fills],
        'fill_qty': [f.signed_qty for f in fills]
    })

    fill_out = dataframe_feeder(df)
    fill_price = fill_out[1].named_ref('fill_price')
    fill_qty = fill_out[2].named_ref('fill_qty')
    fill_event = fill_out[3]

    with graph.control_if(fill_event):
      with graph.control_if(ts >= dump_from_ts and ts < dump_to_ts):
        aggregator = scalar_value_aggregator([
            ts,
            *ask0_past,
            ask0,
            *ask0_fut,
            *bid0_past,
            bid0,
            *bid0_fut,
            *mid_past,
            mid,
            *mid_fut,
            fill_price,
            fill_qty
        ])

  return aggregator


def timed_y_graph(symbol, machine, fills, timed_y_points, date):
  symbol_info = get_symbol_info(symbol, date)
  dump_from_ts, dump_to_ts = get_date_timestamp_range(date)
  fills = [elem for elem in fills if dump_from_ts <= elem.ts < dump_to_ts]
  return timed_y(symbol_info, machine, fills, dump_from_ts, dump_to_ts, timed_y_points)


def generate_timed_y_dataframe(symbol, machine, fills, timed_y_points=None):
  timestamps = [f.ts for f in fills]
  from_ts = min(timestamps)
  to_ts = max(timestamps)

  timed_y_points = timed_y_points or DEFAULT_TIMED_Y_POINTS
  margin_duration = max([to_duration_int(point) for point in timed_y_points]) * 2
  logging.info('Margin duration: %s', to_duration_str(margin_duration))

  df = run_from_fastfeed_multiprocess(functools.partial(timed_y_graph,
                                                        symbol,
                                                        machine,
                                                        fills,
                                                        timed_y_points),
                                      from_ts,
                                      to_ts,
                                      machine,
                                      inject_date_to_graph_func=True,
                                      dump_pending_row=True,
                                      duration_before=margin_duration,
                                      duration_after=margin_duration,
                                      max_workers=None)
  return df
